package com.ai.fortune_matrix.service.impl;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.HealthStatus;
import co.elastic.clients.elasticsearch._types.mapping.*;
import co.elastic.clients.elasticsearch.cluster.HealthResponse;
import co.elastic.clients.elasticsearch.indices.CreateIndexResponse;
import co.elastic.clients.elasticsearch.indices.IndexSettings;


import com.alibaba.cloud.ai.model.RerankModel;
import io.micrometer.common.util.StringUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.autoconfigure.vectorstore.elasticsearch.ElasticsearchVectorStoreProperties;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
/**
 * @Author: zjg
 * @Date: 2025/5/16 17:11
 * @Description:
 **/
@Service
@Slf4j
public class ElasticSearchStore {

    public static final String textField = "content";

    private static final String vectorField = "embedding";

    @Value("classpath:/rag/winning_number_pool.txt")
    private Resource springAiResource;



    private final VectorStore vectorStore;

    private final ChatModel chatModel;


    private final RerankModel rerankModel;

    private final ElasticsearchClient elasticsearchClient;

    private final ElasticsearchVectorStoreProperties elasticsearchVectorStoreProperties;

    public ElasticSearchStore(ChatModel chatModel, VectorStore vectorStore, RerankModel rerankModel,
                              ElasticsearchClient elasticsearchClient, ElasticsearchVectorStoreProperties elasticsearchVectorStoreProperties) {
        this.chatModel = chatModel;
        this.vectorStore = vectorStore;
        this.rerankModel = rerankModel;
        this.elasticsearchClient = elasticsearchClient;
        this.elasticsearchVectorStoreProperties = elasticsearchVectorStoreProperties;
    }
    public void importDocuments() {
        // 1. 检查资源文件是否存在
        if (!springAiResource.exists()) {
            throw new IllegalStateException("Resource file not found: " + springAiResource.getFilename());
        }

        // 2. 创建索引（如果不存在）
        createIndexIfNotExists();

        // 1. parse document
        TextReader reader = new TextReader(springAiResource);
        List<Document> documents = reader.get();
        log.info("{} documents loaded", documents.size());

        // 2. split trunks
        List<Document> splitDocuments = new TokenTextSplitter().apply(documents);
        log.info("{} documents split", splitDocuments.size());

        // 3. create embedding and store to vector store
        log.info("create embedding and save to vector store");
        vectorStore.add(splitDocuments);
        try {
            elasticsearchClient.indices().refresh(r -> r.index(elasticsearchVectorStoreProperties.getIndexName()));
        } catch (IOException e) {
            throw new RuntimeException("Failed to import documents", e);
        }
        verifyDataImported();
    }

    private void createIndexIfNotExists() {
        try {
            String indexName = elasticsearchVectorStoreProperties.getIndexName();
            Integer dimsLength = elasticsearchVectorStoreProperties.getDimensions();

            if (StringUtils.isBlank(indexName)) {
                throw new IllegalArgumentException("Elastic search index name must be provided");
            }

            boolean exists = elasticsearchClient.indices().exists(idx -> idx.index(indexName)).value();
            if (exists) {
                log.debug("Index {} already exists. Skipping creation.", indexName);
                return;
            }

            String similarityAlgo = elasticsearchVectorStoreProperties.getSimilarity().name();
            IndexSettings indexSettings = IndexSettings
                    .of(settings -> settings.numberOfShards(String.valueOf(1)).numberOfReplicas(String.valueOf(1)));

            // Maybe using json directly?
            Map<String, Property> properties = new HashMap<>();
            properties.put(vectorField, Property.of(property -> property.denseVector(
                    DenseVectorProperty.of(dense -> dense.index(true).dims(dimsLength).similarity(similarityAlgo)))));
            properties.put(textField, Property.of(property -> property.text(TextProperty.of(t -> t))));

            Map<String, Property> metadata = new HashMap<>();
            metadata.put("ref_doc_id", Property.of(property -> property.keyword(KeywordProperty.of(k -> k))));

            properties.put("metadata",
                    Property.of(property -> property.object(ObjectProperty.of(op -> op.properties(metadata)))));

            CreateIndexResponse indexResponse = elasticsearchClient.indices()
                    .create(createIndexBuilder -> createIndexBuilder.index(indexName)
                            .settings(indexSettings)
                            .mappings(TypeMapping.of(mappings -> mappings.properties(properties))));

            if (!indexResponse.acknowledged()) {
                throw new RuntimeException("failed to create index");
            }

            log.info("create elasticsearch index {} successfully", indexName);
        }
        catch (IOException e) {
            log.error("failed to create index", e);
            throw new RuntimeException(e);
        }
    }

    private void verifyDataImported() {
        String indexName = elasticsearchVectorStoreProperties.getIndexName();

        // 检查索引状态
        HealthResponse health = null;
        try {
            health = elasticsearchClient.cluster().health(b -> b.index(indexName));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        if (health.status() == HealthStatus.Red) {
            throw new IllegalStateException("Index is in RED state");
        }

        // 查询文档数量
        long count = 0;
        try {
            count = elasticsearchClient.count(b -> b.index(indexName)).count();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        log.info("Index {} currently contains {} documents", indexName, count);
    }

}
